{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9128b19e-b704-45b8-97fb-0f2355a7630e",
   "metadata": {},
   "source": [
    "# Processing Experimental Data\n",
    "\n",
    "This file takes the raw experimental data at the four signal levels we studied and produces approximates for the training and testing portions of the datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dce1e009",
   "metadata": {},
   "outputs": [],
   "source": [
    "import h5py\n",
    "import numpy as np\n",
    "import torch as t\n",
    "from RPI_tools import pytorch_tools\n",
    "from tqdm import tqdm\n",
    "from scipy import io\n",
    "\n",
    "# Edit this line to point to the location of the data folder on your system.\n",
    "data_folder_prefix = 'data/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97864223",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "photon_levels = [1, 10, 100, 1000]\n",
    "for photon in photon_levels:\n",
    "    \n",
    "    print('Processing data for target photon level %d.' % photon)\n",
    "\n",
    "    # We load the raw patterns from the experimental datasets\n",
    "    filename = data_folder_prefix + ('Experimental/RPI_Datasets/Collection_3_photon_%d.npy' % photon)\n",
    "    \n",
    "    patterns = np.load(filename)\n",
    "    print(np.shape(patterns))\n",
    "    \n",
    "    # The first 4,000 images are used as the training data\n",
    "    train = patterns[:4000]\n",
    "    # And the next 100 images are the test dataset.\n",
    "    test = patterns[4000:4100]\n",
    "    \n",
    "    # Now, we load the appropriate probe function, as produced from a separate ptychography scan.\n",
    "    calibration_data = io.loadmat(data_folder_prefix + \n",
    "                                  ('Experimental/Ptychography_Calibrations/Calibrated_Ptychography_Result_EMGain3800_%dPhotonTarget.mat' % photon))\n",
    "    expanded_probe = calibration_data['probe'][0]\n",
    "    background = calibration_data['background']\n",
    "\n",
    "    print('Dataset and calibration data loaded.')\n",
    "    \n",
    "    lr = 1\n",
    "    iters = 1\n",
    "\n",
    "    # The iterative algorithm is defined in pytorch for historical reasons. This is the reason that\n",
    "    # this codebase depends on both pytorch and tensorflow.\n",
    "    expanded_probe0 = t.from_numpy(expanded_probe)\n",
    "    background = t.from_numpy(background)\n",
    "    train_patterns = t.from_numpy(train)\n",
    "    test_patterns = t.from_numpy(test)\n",
    "\n",
    "    # Here, we run one iteration of the automatic-differentiation RPI algorithm to produce the training approximants\n",
    "    train_approx = []\n",
    "    for pattern in tqdm(train_patterns):\n",
    "        result, _ = pytorch_tools.reconstruct(pattern, expanded_probe0, 256, lr, iters, background=background, \n",
    "                                              loss_func=pytorch_tools.amplitude_mse, GPU=True, saturation=2**14)\n",
    "        # This accounts for a 90 degree rotation between the definition of the images and the detector orientation\n",
    "        train_approx.append(np.rot90(result.numpy(), 2))\n",
    "        \n",
    "    print('Training approximants produced.')\n",
    "\n",
    "    # Here, we run one iteration of the automatic-differentiation RPI algorithm to produce the testing approximants\n",
    "    test_approx = []\n",
    "    for pattern in test_patterns:\n",
    "        result, _ = pytorch_tools.reconstruct(pattern, expanded_probe0, 256, lr, iters, background=background, \n",
    "                                              loss_func=pytorch_tools.amplitude_mse, GPU=True, saturation=2**14)\n",
    "        # This accounts for a 90 degree rotation between the definition of the images and the detector orientation\n",
    "        test_approx.append(np.rot90(result.numpy(), 2))\n",
    "\n",
    "    print('Testing approximants produced.')\n",
    "    \n",
    "    # The approximant is defined to just be the phase channel of the output.\n",
    "    train_approx = np.angle(np.array(train_approx))\n",
    "    test_approx = np.angle(np.array(test_approx))\n",
    "\n",
    "    # Finally, we save out the approximants that we've created\n",
    "    np.save('train-approx-%d-iter-%d-lr-%0.2f.npy' % (photon, iters, lr), train_approx)\n",
    "    np.save('test-approx-%d-iter-%d-lr-%0.2f.npy' % (photon, iters, lr), test_approx)\n",
    "    \n",
    "print('All approximants saved, have a nice day!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d56784b-942d-4af9-9358-af7682a5de42",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "Xs, Ys = np.mgrid[:256,:256]\n",
    "Xs = Xs - np.mean(Xs)\n",
    "Ys = Ys - np.mean(Ys)\n",
    "Rs = np.sqrt(Xs**2 + Ys**2)\n",
    "\n",
    "photon_levels = [1, 10, 100, 1000]\n",
    "for photon in photon_levels:\n",
    "    test_approx = np.load('test-approx-%d-iter-%d-lr-%0.2f.npy' % (photon, iters, lr))\n",
    "    test_approx[:, Rs>128] = 0\n",
    "    \n",
    "    plt.rcParams['figure.figsize'] = [20, 20]\n",
    "    fig = plt.figure()\n",
    "    fig.subplots_adjust(hspace=.05, wspace=.05)\n",
    "\n",
    "    for i in range(16):\n",
    "        ax = fig.add_subplot(4, 4, i +1)\n",
    "        ax.axes.xaxis.set_visible(False)\n",
    "        ax.axes.yaxis.set_visible(False)\n",
    "        plt.imshow(test_approx[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5dcb7565",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get 100 iteration reconstruction results\n",
    "\n",
    "photon_levels = [1, 10, 100, 1000]\n",
    "for photon in photon_levels:\n",
    "    \n",
    "    print('Processing data for target photon level %d.' % photon)\n",
    "\n",
    "    # We load the raw patterns from the experimental datasets\n",
    "    filename = data_folder_prefix + ('Experimental/RPI_Datasets/Collection_3_photon_%d.npy' % photon)\n",
    "    \n",
    "    patterns = np.load(filename)\n",
    "    print(np.shape(patterns))\n",
    "    \n",
    "    # get the test dataset.\n",
    "    test = patterns[4000:4100]\n",
    "    \n",
    "    # Now, we load the appropriate probe function, as produced from a separate ptychography scan.\n",
    "    calibration_data = io.loadmat(data_folder_prefix + \n",
    "                                  ('Experimental/Ptychography_Calibrations/Calibrated_Ptychography_Result_EMGain3800_%dPhotonTarget.mat' % photon))\n",
    "    expanded_probe = calibration_data['probe'][0]\n",
    "    background = calibration_data['background']\n",
    "\n",
    "    print('Dataset and calibration data loaded.')\n",
    "    \n",
    "    lr = 0.5\n",
    "    iters = 1000\n",
    "\n",
    "    # The iterative algorithm is defined in pytorch for historical reasons. This is the reason that\n",
    "    # this codebase depends on both pytorch and tensorflow.\n",
    "    expanded_probe0 = t.from_numpy(expanded_probe)\n",
    "    background = t.from_numpy(background)\n",
    "    test_patterns = t.from_numpy(test)\n",
    "\n",
    "    # Here, we run one iteration of the automatic-differentiation RPI algorithm to produce the testing approximants\n",
    "    test_approx = []\n",
    "    for pattern in test_patterns:\n",
    "        \n",
    "        result, _ = pytorch_tools.reconstruct(pattern, expanded_probe0, 256, lr, iters, background=background, \n",
    "                                                  loss_func=pytorch_tools.amplitude_mse, GPU=True, saturation=2**14)\n",
    "        # This accounts for a 90 degree rotation between the definition of the images and the detector orientation\n",
    "        test_approx.append(np.rot90(result.numpy(), 2))\n",
    "\n",
    "    print('Testing iterative reconstructions produced.')\n",
    "    \n",
    "    # The approximant is defined to just be the phase channel of the output.\n",
    "    test_approx = np.angle(np.array(test_approx))\n",
    "\n",
    "    # Finally, we save out the approximants that we've created\n",
    "    np.save('test-approx-%d-iter-%d-lr-%0.2f.npy' % (photon, iters, lr), test_approx)\n",
    "    \n",
    "print('All approximants saved, have a nice day!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78434ac5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "\n",
    "Xs, Ys = np.mgrid[:256,:256]\n",
    "Xs = Xs - np.mean(Xs)\n",
    "Ys = Ys - np.mean(Ys)\n",
    "Rs = np.sqrt(Xs**2 + Ys**2)\n",
    "\n",
    "photon_levels = [1, 10, 100, 1000]\n",
    "for photon in photon_levels:\n",
    "    \n",
    "    test_approx = np.load('test-approx-%d-iter-%d-lr-%0.2f.npy' % (photon, iters, lr))\n",
    "    test_approx[:, Rs>128] = 0\n",
    "    \n",
    "#     if photon == 1e3:\n",
    "#         test_approx = np.load('test-approx-%d-iter-%d-lr-%0.2f.npy' % (photon, iters, 1))\n",
    "#         test_approx[:, Rs>128] = 0\n",
    "    \n",
    "    plt.rcParams['figure.figsize'] = [20, 20]\n",
    "    fig = plt.figure()\n",
    "    fig.subplots_adjust(hspace=.05, wspace=.05)\n",
    "\n",
    "    for i in range(16):\n",
    "        ax = fig.add_subplot(4, 4, i +1)\n",
    "        ax.axes.xaxis.set_visible(False)\n",
    "        ax.axes.yaxis.set_visible(False)\n",
    "        plt.imshow(test_approx[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38ee15e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "R = 0.5\n",
    "\n",
    "photon_levels = [1, 10, 100, 1000]\n",
    "for photon in photon_levels:\n",
    "    \n",
    "    print('Processing E2E data for target photon level %d.' % photon)\n",
    "\n",
    "    # We load the raw patterns from the experimental datasets\n",
    "    filename = data_folder_prefix + ('Experimental/RPI_Datasets/Collection_3_photon_%d.npy' % photon)\n",
    "    \n",
    "    patterns = np.load(filename)\n",
    "    print(np.shape(patterns))\n",
    "    \n",
    "    # The first 4,000 images are used as the training data\n",
    "    train = patterns[:4000]\n",
    "    # And the next 100 images are the test dataset.\n",
    "    test = patterns[4000:4100]\n",
    "    \n",
    "    # Now, we load the appropriate probe function, as produced from a separate ptychography scan.\n",
    "    calibration_data = io.loadmat(data_folder_prefix + \n",
    "                                  ('Experimental/Ptychography_Calibrations/Calibrated_Ptychography_Result_EMGain3800_%dPhotonTarget.mat' % photon))\n",
    "    expanded_probe = calibration_data['probe'][0]\n",
    "    background = calibration_data['background']\n",
    "    \n",
    "    # remove the background for experimental data\n",
    "    tr_patterns = []\n",
    "    for img in train:\n",
    "        img -= background\n",
    "        tr_patterns.append(img)\n",
    "    tr_patterns = np.array(tr_patterns)\n",
    "    # zero pad the data to (4000, 1024, 1024)\n",
    "    tr_patterns = np.pad(tr_patterns, ((0, 0), (16, 16), (16, 16)))\n",
    "    \n",
    "    # remove the background for experimental data\n",
    "    test_patterns = []\n",
    "    for img in train:\n",
    "        img -= background\n",
    "        test_patterns.append(img)\n",
    "    test_patterns = np.array(test_patterns)\n",
    "    # zero pad the data to (100, 1024, 1024)\n",
    "    test_patterns = np.pad(test_patterns, ((0, 0), (16, 16), (16, 16)))\n",
    "    \n",
    "    probe = np.pad(expanded_probe, ((16, 16), (16, 16)))\n",
    "    \n",
    "    chn = 4\n",
    "    \n",
    "    flatten = np.zeros((4000, 256, 256, chn**2)).astype(np.float32)\n",
    "\n",
    "    for i in range(4000):\n",
    "        for j in range(chn):\n",
    "            for k in range(chn):\n",
    "                flatten[i, :, :, chn*j + k] = tr_patterns[i, 256*j:256*(j+1), 256*k:256*(k+1)]\n",
    "                \n",
    "    flatten = flatten.astype(np.float32)\n",
    "\n",
    "    np.save('exp-tr_patterns-flatten-R-%0.2f-phperpix-%d.npy' % (R, photon), flatten)\n",
    "    \n",
    "    \n",
    "    flatten = np.zeros((100, 256, 256, chn**2))\n",
    "\n",
    "    for i in range(100):\n",
    "        for j in range(chn):\n",
    "            for k in range(chn):\n",
    "                flatten[i, :, :, chn*j + k] = test_patterns[i, 256*j:256*(j+1), 256*k:256*(k+1)]\n",
    "                \n",
    "    flatten = flatten.astype(np.float32)\n",
    "\n",
    "    np.save('exp-test_patterns-flatten-R-%0.2f-phperpix-%d.npy' % (R, photon), flatten)\n",
    "    \n",
    "\n",
    "    flatten = np.zeros((256, 256, 2*chn**2))\n",
    "\n",
    "    for j in range(chn):\n",
    "        for k in range(chn):\n",
    "            flatten[:, :, chn*j + k] = np.abs(probe[256*j:256*(j+1), 256*k:256*(k+1)])\n",
    "            flatten[:, :, chn**2 + chn*j + k] = np.angle(probe[256*j:256*(j+1), 256*k:256*(k+1)])\n",
    "            \n",
    "    flatten = flatten.astype(np.float32)\n",
    "\n",
    "    np.save('exp-probe-flatten-R-%0.2f-phperpix-%d.npy' % (R, photon), flatten)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2be42b7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
